/*
 * Copyright 2003 Program of Computer Graphics, Cornell University
 *     580 Rhodes Hall
 *     Cornell University
 *     Ithaca NY 14853
 * Web: http://www.graphics.cornell.edu/
 * 
 * Not for commercial use. Do not redistribute without permission.
 */

//This file contains all the routines needed during KDTree traversal.
//They are all gathered here in one source file so that the compiler will have
//the easiest time optimizing them as a group
#include "stdafx.h"

#include "KDTree.h"
#include "Ray.h"
#include "PrimitiveList.h"
#include "Triangle.h"
#include "BriteDefs.h"
#include <float.h> //    FLT_MIN, FLT_MAX, FLT_EPSILON

//recorded state about a KDNode used during tree traveral for ray intersection
//includes ray parameters for entry and exit from the nodes bounding cell
class KDState {
public:
  float enterT, exitT;
  class KDNodeHandle *nodeHandle;
  KDState *nextValid;
};

class ALIGN16 KDQuadState {
public:
  F32vec4 enterT, exitT;
  class KDNodeHandle *nodeHandle;
//  KDQuadState *nextValid;
  int active;
};

//take a four bit intex and return zero if it is zero and -1 otherwise
//static signed char bitSet4Table[16] = { 0, -1, -1, -1, -1, -1, -1, -1,
//                                       -1, -1, -1, -1, -1, -1, -1, -1};

//enable slower but more accurate methods that use double precision intermediate calculations
//use doubles when computing intersections with kd-tree split planes
//#define DOUBLE_SPLIT_TRAVERSE
//use doubles when computing ray-triangle interesections
//#define DOUBLE_TRIANGLE_INTERSECT
//disallow early termination of kd-traversal when cells are beyond the end of the ray
//#define DISABLE_EARLY_TRAVERSAL_TERMINATION

#ifndef DOUBLE_TRIANGLE_INTERSECT
inline static F32vec4 lowerSum(const F32vec4 &sum, const F32vec4 &a, const F32vec4 &b) {
  F32vec4 bv = sum - a;
  return (a - (sum - bv)) + (b - bv);
}
inline static F32vec4 lowerDiff(const F32vec4 &diff, const F32vec4 &a, const F32vec4 &b) {
  F32vec4 bv = diff - a;
  return (a - (diff - bv)) - (b + bv);
}

inline static void unpack4(F32vec4 &x, F32vec4 &y, F32vec4 &z, const F32vec4 &v0,
                           const F32vec4 &v1, const F32vec4 &v2, const F32vec4 &v3) {
  F32vec4 xxyy01 = _mm_unpacklo_ps(v0,v1);
  F32vec4 xxyy23 = _mm_unpacklo_ps(v2,v3);
  x = _mm_movelh_ps(xxyy01,xxyy23);
  y = _mm_movehl_ps(xxyy23,xxyy01);
  z = _mm_movelh_ps(_mm_unpackhi_ps(v0,v1),_mm_unpackhi_ps(v2,v3));
}
//version which improves numerical precision significantly by computing a vector
// P = S - aD using extended precision arithmetic in order to transform the into
//a more robust one where P is nearly orthogonal to D
static void intersectQuadPackedTriangleList(QuadPackedTriangleList *list, RayExt *ray, 
                                            IntersectionRecord *hitRec)
{
  F32vec4 tStart = F32vec4(ray->tStart);
  F32vec4 tEnd = F32vec4(ray->tEnd);
  const Is32vec4 highMaskInt(0xFFFFF000,0xFFFFF000,0xFFFFF000,0xFFFFF000);
  const F32vec4 highMask = _mm_castsi128_ps(highMaskInt);
  assert(tStart[0]<=tEnd[0]);
  {
//_asm int 3; //force debugger breakpoint
    QuadPackedTriangles *tris = list->tris;
#ifdef USE_COMPILED_SHARED_TRIANGLE
    for(;;) {  //will break when PL_QUADPACKEDTRIANGLE_END flag is found
      F32vec4 e1x,e1y,e1z,e2x,e2y,e2z,v0x,v0y,v0z;
      {
      CompiledSharedTriangle *t0 = tris->getTri0Ptr();
      CompiledSharedTriangle *t1 = tris->tri[1];
      CompiledSharedTriangle *t2 = tris->tri[2];
      CompiledSharedTriangle *t3 = tris->tri[3];
      unpack4(e1x,e1y,e1z,t0->e1,t1->e1,t2->e1,t3->e1);
      unpack4(e2x,e2y,e2z,t0->e2,t1->e2,t2->e2,t3->e2);
      unpack4(v0x,v0y,v0z,t0->v0,t1->v0,t2->v0,t3->v0);
      }
#else //USE_COMPILED_SHARED_TRIANGLE
    for(int count = list->length;count>0;count-=4) {
      F32vec4 e1x = tris->e1x;
      F32vec4 e1y = tris->e1y;
      F32vec4 e1z = tris->e1z;
      F32vec4 e2x = tris->e2x;
      F32vec4 e2y = tris->e2y;
      F32vec4 e2z = tris->e2z;
      F32vec4 v0x = tris->v0x;
      F32vec4 v0y = tris->v0y;
      F32vec4 v0z = tris->v0z;
#endif //USE_COMPILED_SHARED_TRIANGLE
      //C is e1 X e2
      F32vec4 Cx = e1y*e2z - e1z*e2y;
      F32vec4 Cy = e1z*e2x - e1x*e2z;
      F32vec4 Cz = e1x*e2y - e1y*e2x;
//      F32vec4 denom = F32vec4(1.0f)/(ray->v_dirX*Cx + ray->v_dirY*Cy + ray->v_dirZ*Cz);
      F32vec4 denom = rcp_nr(ray->v_dirX*Cx + ray->v_dirY*Cy + ray->v_dirZ*Cz);
      //S is vertex zero minus the ray origin
      F32vec4 Sx = v0x - ray->v_originX;
      F32vec4 Sy = v0y - ray->v_originY;
      F32vec4 Sz = v0z - ray->v_originZ;
      //next we compute a "mostly" orthogonalized version of S by subtracting
      //a scaled version of D (ray direction) from S to form the vector
      //P = S - SDu * D, that is much closer to being orthogonal to D.
      //SDu is S dot D except truncated to 12 bits of precision
      F32vec4 SDu = (Sx*ray->v_dirX + Sy*ray->v_dirY + Sz*ray->v_dirZ)&highMask;
      //we then use extended precision arithmetic (nearly double precision) to
      //form P = S - SDu * D to reduce roundoff error
      //D = Du + Dl where Du and Dl each have at most 12 significant bits
      F32vec4 Dux = ray->v_dirX&highMask;
      F32vec4 Dlx = ray->v_dirX - Dux;
      F32vec4 Duy = ray->v_dirY&highMask;
      F32vec4 Dly = ray->v_dirY - Duy;
      F32vec4 Duz = ray->v_dirZ&highMask;
      F32vec4 Dlz = ray->v_dirZ - Duz;
      //J = Ju + Jl = SDu * D using exact arithmetic
      F32vec4 Jux = SDu*ray->v_dirX;
      F32vec4 Jlx = SDu*Dlx - (Jux - SDu*Dux);
      F32vec4 Juy = SDu*ray->v_dirY;
      F32vec4 Jly = SDu*Dly - (Juy - SDu*Duy);
      F32vec4 Juz = SDu*ray->v_dirZ;
      F32vec4 Jlz = SDu*Dlz - (Juz - SDu*Duz);
      //K = Ku + Kl = J + O (ray origin)
      F32vec4 Kux = Jux + ray->v_originX;
      F32vec4 Klx = Jlx + lowerSum(Kux,Jux,ray->v_originX);
      F32vec4 Kuy = Juy + ray->v_originY;
      F32vec4 Kly = Jly + lowerSum(Kuy,Juy,ray->v_originY);
      F32vec4 Kuz = Juz + ray->v_originZ;
      F32vec4 Klz = Jlz + lowerSum(Kuz,Juz,ray->v_originZ);
      //P = Pu + Pl = V0 - K (triangle vertex zero)
/*      F32vec4 Pux = tris->v0x - Kux;
      F32vec4 Plx = lowerDiff(Pux,tris->v0x,Kux) - Klx;
      F32vec4 Puy = tris->v0y - Kuy;
      F32vec4 Ply = lowerDiff(Puy,tris->v0y,Kuy) - Kly;
      F32vec4 Puz = tris->v0z - Kuz;
      F32vec4 Plz = lowerDiff(Puz,tris->v0z,Kuz) - Klz;
      //next we form the P which is the component of S which is "nearly" perpendicular to D
      //the idea is that B = S X D = P X D, but P X D has less roundoff error (better conditioned)
      //we assume that D (the ray direction) is normalized so that D dot D = 1
      F32vec4 Px = Pux + Plx;
      F32vec4 Py = Puy + Ply;
      F32vec4 Pz = Puz + Plz;
      */
      F32vec4 Px = (v0x - Kux) - Klx;
      F32vec4 Py = (v0y - Kuy) - Kly;
      F32vec4 Pz = (v0z - Kuz) - Klz;
      //t is the parametric parameter along the ray of the intersection
      //t = (S dot C) / (D dot C)
//      F32vec4 t = (Sx*Cx + Sy*Cy + Sz*Cz)*denom;
      //or better t = SDu + ((P dot C) / (D dot C))
      F32vec4 t = SDu + (Px*Cx + Py*Cy + Pz*Cz)*denom;
      //B is P X D = S X D
      F32vec4 Bx = Py*ray->v_dirZ - Pz*ray->v_dirY;
      F32vec4 By = Pz*ray->v_dirX - Px*ray->v_dirZ;
      F32vec4 Bz = Px*ray->v_dirY - Py*ray->v_dirX;
      //u is the first barycentric coordinate of the intersection
      //u = (e2 dot B) / (D dot C)
      F32vec4 u = (e2x*Bx + e2y*By + e2z*Bz)*denom;
      //v is the second barycentric coorinate of the intersection
      //v = -(e1 dot B) / (D dot C)
      //we compute negV = -v instead to avoid the extra operation
      F32vec4 negV = (e1x*Bx + e1y*By + e1z*Bz)*denom;
      //test to see if any intersections are valid
      F32vec4 v_zero = _mm_setzero_ps();
      F32vec4 v_one = F32vec4(1.0f);
      F32vec4 valid = cmpge(u,v_zero) & cmple(negV,v_zero) & cmple(u-negV,v_one)
        & cmpgt(t,tStart) & cmplt(t,tEnd);
      int test = move_mask(valid);
//      prefetch(((char *)tris)+256);  //in case we have more triangles in this list
      if (test) { //there is at least one valid intersection
        if (hitRec->func) { //if we have a custom intersection responder, then use it
          int bit = 0x01;
          for(int i=0;i<4;i++) {
            if (test&bit) {
              (*hitRec->func)(ray,hitRec,t[i],u[i],-negV[i],0,tris->getId(i));
            }
            bit = bit+bit;
          }
          tEnd = F32vec4(minv(ray->tEnd,tEnd[0]));
        } else { //otherwise the default is find first intersection
//_asm int 3; //force debugger breakpoint
          int bit = 0x01;
          for(int i=0;i<4;i++) {
            if ((test&bit)&&(t[i]<tEnd[0])) {
              hitRec->t = ray->tEnd = t[i];
              hitRec->u = u[i];
              hitRec->v = -negV[i];
              hitRec->w = 0.0f;
              hitRec->primitiveId = tris->getId(i);
              tEnd = F32vec4(t[i]);
            }
            bit = bit+bit;
          }
        }
      }
#ifdef USE_COMPILED_SHARED_TRIANGLE
      if (tris->isLastQuadPack()) break;
#endif //USE_COMPILED_SHARED_TRIANGLE
      tris++;
    }
  }
}

#else //DOUBLE_TRIANGLE_INTERSECT
//This version uses doubles for intermediate quantities during triangle intersection.
//As such its slower but higher precision than the faster method above that uses floats and SSE
static void intersectQuadPackedTriangleList(QuadPackedTriangleList *list, RayExt *ray, 
                                            IntersectionRecord *hitRec)
{
  double tStart = ray->tStart;
  double tEnd = ray->tEnd;
  assert(tStart<=tEnd);
  {
_asm int 3; //force debugger breakpoint
    QuadPackedTriangles *tris = list->tris;
    int count = list->length;
    while(count > 0) {
      for(int i=0;(i<4)&&(count>0);i++,count--) {
        //first form edges
        double e1x = tris->e1x[i];
        double e1y = tris->e1y[i];
        double e1z = tris->e1z[i];
        double e2x = tris->e2x[i];
        double e2y = tris->e2y[i];
        double e2z = tris->e2z[i];
        //C is e1 X e2
        double Cx = e1y*e2z - e1z*e2y;
        double Cy = e1z*e2x - e1x*e2z;
        double Cz = e1x*e2y - e1y*e2x;
        double denom = 1.0/(Cx*ray->dirX() + Cy*ray->dirY() + Cz*ray->dirZ());
        //denom is 1.0 / dot(C,ray.direction)
        //S is vertex zero minus the ray origin
//        double Sx = tris->v0x[i] - (double)ray->originX();
//        double Sy = tris->v0y[i] - (double)ray->originY();
//        double Sz = tris->v0z[i] - (double)ray->originZ();
        float Sx = tris->v0x[i] - ray->originX();
        float Sy = tris->v0y[i] - ray->originY();
        float Sz = tris->v0z[i] - ray->originZ();
        //t is the parametric parameter along the ray of the intersection
        //t = (S dot C) / (D dot C)
        double t = (Sx*Cx + Sy*Cy + Sz*Cz)*denom;
        //B is S X D
        double Bx = Sy*(double)ray->dirZ() - Sz*(double)ray->dirY();
        double By = Sz*(double)ray->dirX() - Sx*(double)ray->dirZ();
        double Bz = Sx*(double)ray->dirY() - Sy*(double)ray->dirX();
        //u is the first barycentric coordinate of the intersection
        //u = (e2 dot B) / (D dot C)
        double u = (e2x*Bx + e2y*By + e2z*Bz)*denom;
        //v is the second barycentric coorinate of the intersection
        //v = -(e1 dot B) / (D dot C)
        double v = -(e1x*Bx + e1y*By + e1z*Bz)*denom;
        //test to see if intersection is valid
        if ((t>=tStart)&&(t<=tEnd)&&(u>=0)&&(v>=0)&&(u+v<=1.0)) {
          if (hitRec->func) { //if we have a custom intersection responder, then use it
            (*hitRec->func)(ray,hitRec,t,u,v,0,tris->id[i]);
            tEnd = minv(tEnd,(double)ray->tEnd); //reload in case ray was shortened
          } else { //otherwise the default is find first intersection
//_asm int 3; //force debugger breakpoint
            hitRec->t = ray->tEnd = tEnd =  t;
            hitRec->primitiveId = tris->id[i];
            hitRec->u = u;
            hitRec->v = v;
            hitRec->w = 0;
          }
        }
      }
      tris++;
    }
  }
}
#endif //DOUBLE_TRIANGLE_INTERSECT

static inline LONG_PTR Select(LONG_PTR a, LONG_PTR b, LONG_PTR mask) {
//  return (a&mask) | (b&~mask);
  return b ^ (mask & (a ^ b));
}
void KDTree::intersect(const Ray &ray, IntersectionRecord *hitRec)
{
  KDState *cur;
  RayExt extRay;
  KDState stack[KDTREE_MAX_DEPTH+2];
  extRay.initializeFromRay(ray);
  if (!epsBounds.clipRayExt(extRay)) {
    return;
  }
  //stack[0] is a dummy node that we use to detect when we finish.  It should always be past the end of the ray
  stack[0].enterT = extRay.tEnd + fabsf(extRay.tEnd) + 1.0f;
  stack[0].exitT = stack[0].enterT;
  stack[0].nodeHandle = makeKDNodeLeafHandle((PrimitiveList*)&dummy);
  stack[0].nextValid = &stack[0];
  stack[1].enterT = extRay.tStart;
  stack[1].exitT = extRay.tEnd;
//  stack[1].enterT = -FLT_MAX;
//  stack[1].exitT = FLT_MAX;
  stack[1].nodeHandle = rootHandle;
  stack[1].nextValid = &stack[1];
  cur = &stack[1];
  
  //traverse the tree until we find a cell past the end of the clipped ray
  F32vec1 rayEndT(extRay.tEnd);
  for(;;) {
    const KDNode *node = getNodePtr(cur->nodeHandle);
    const int type = getNodeType(cur->nodeHandle);

    if (type == KDST_LEAF) {
#ifndef DISABLE_EARLY_TRAVERSAL_TERMINATION
      if (_mm_ucomigt_ss(F32vec1(cur->enterT),rayEndT)) { return; } //once we reach a cell past the end of the ray, we are done
#else
      if (cur == &stack[0]) { return; } //reached end of traversal of the entire tree
#endif
      PrimitiveList *list = (PrimitiveList*)node;
      while(list!=NULL) {
        if (list->isQuadPackedTriangleList()) {
          intersectQuadPackedTriangleList((QuadPackedTriangleList*)list,&extRay,hitRec);
          break; //these must appear last in the linked list of primitives
        }
        (*list->func)(list,&extRay,hitRec);
        list = list->next;
      }
      rayEndT = F32vec1(extRay.tEnd); //reload in case it was changed
      //warning: if we use, rayEndT = extRay.tEnd, the compiler truncates value to int for some strange reason
      cur = cur[-1].nextValid;
      continue;
    }
    //otherwise it must be a splitting node
#ifndef DOUBLE_SPLIT_TRAVERSE
    const F32vec1 splitT = F32vec1((node->splitValue - extRay.originComponent(type))
                              * extRay.invDirComponent(type));
#else
    //double version of this code
    const F32vec1 splitT = F32vec1((node->splitValue - (double)extRay.originComponent(type))
                                   / extRay.direction[type]);
#endif
    const F32vec1 enterT = F32vec1(cur->enterT);
    const F32vec1 exitT = F32vec1(cur->exitT);
    int leftIndex = extRay.dirType(type);
    cur[leftIndex].nodeHandle = node->left;
    cur[1-leftIndex].nodeHandle = node->right;
    LONG_PTR cur0Valid = _mm_ucomile_ss(splitT,exitT)?-1:0;
    LONG_PTR cur1Valid = _mm_ucomile_ss(enterT,splitT)?-1:0;
//    int cur0Valid = move_mask(F32vec4(cmpge(splitT,exitT)));
//    int cur1Valid = move_mask(F32vec4(cmpge(enterT,splitT)));
    Store_F32vec1(&cur[1].enterT,enterT);
    Store_F32vec1(&cur[1].exitT,simd_min(exitT,splitT));
    Store_F32vec1(&cur[0].enterT,simd_max(splitT,enterT));
//    cur0Valid = (cur0Valid&0x01)-1;
    //if it is an empty leaf cell, then we can skip it (same as if ray missed the cell)
    if (KDST_LEAF == (LONG_PTR)cur[0].nodeHandle) cur0Valid = 0;
    cur0Valid = Select((LONG_PTR)cur,(LONG_PTR)(cur[-1].nextValid),cur0Valid);
    cur[0].nextValid = (KDState*)cur0Valid;
//    cur1Valid = (cur1Valid&0x01)-1;
    if (KDST_LEAF == (LONG_PTR)cur[1].nodeHandle) cur1Valid = 0;
    cur = (KDState*)Select((LONG_PTR)(&cur[1]),cur0Valid,cur1Valid);
#ifndef DISABLE_EARLY_TRAVERSAL_TERMINATION
    if (_mm_ucomigt_ss(enterT,rayEndT)) { return; } //once we reach a cell past the end of the ray, we are done
#endif
  }
}

void KDTree::intersect(int numRays, const Ray ray[4], IntersectionRecord hitRec[4])
{
  assert((numRays>0)&&(numRays<=4));
  KDQuadState *cur;
  QuadRay quadRay;
  KDQuadState stack[KDTREE_MAX_DEPTH+2];
  int i=0;
  for(;i<numRays;i++) {
    quadRay.ray[i].initializeFromRay(ray[i]);
    epsBounds.clipRayExt(quadRay.ray[i]);
  }
  quadRay.initFromRays(numRays); //setup SIMD vectors for traversal

  //stack[0] is a dummy node that we use to detect when we finish.  It should always be past the end of the ray
  stack[0].enterT = quadRay.quadTEnd + fabs(quadRay.quadTEnd) + F32vec4(1.0f);
  stack[0].exitT = stack[0].enterT;
//  stack[0].nodeHandle = makeKDNodeLeafHandle((PrimitiveList*)&dummy);
  stack[0].nodeHandle = makeKDNodeLeafHandle((PrimitiveList*)NULL);
//  stack[0].nextValid = &stack[0];
  stack[0].active = 0xF;
  stack[1].enterT = quadRay.quadTStart;
  stack[1].exitT = quadRay.quadTEnd;
  stack[1].nodeHandle = rootHandle;
//  stack[1].nextValid = &stack[1];
  stack[1].active = 0xF;
  cur = &stack[2];
  int liveRayMask = move_mask(cmple(quadRay.quadTStart,quadRay.quadTEnd));
  if (!liveRayMask) return;

  //traverse the tree
  for(;;) {
    cur--;
    if (!(cur->active&liveRayMask)) { continue; }

    const KDNode *node = getNodePtr(cur->nodeHandle);
    const int type = getNodeType(cur->nodeHandle);

    if (type == KDST_LEAF) {
      if (node!=NULL) {
        int activeMask = cur->active&liveRayMask;
//        int activeMask = move_mask(cmple(cur->enterT,simd_min(cur->exitT,quadRay.quadTEnd)));
//        if (cur==stack) return;
        int j=0;
        for(; activeMask; activeMask>>=1,j++) {
          if (activeMask&0x01) {
            PrimitiveList *list = (PrimitiveList*)node;
            while(list!=NULL) {
              if (list->isQuadPackedTriangleList()) {
                intersectQuadPackedTriangleList((QuadPackedTriangleList*)list,&quadRay.ray[j],&hitRec[j]);
                break; //these must appear last in the linked list of primitives
              }
              (*list->func)(list,&quadRay.ray[j],&hitRec[j]);
              list = list->next;
            }
            quadRay.quadTEnd[j] = quadRay.ray[j].tEnd; //reload in case it was changed
          }
        }
      }
      liveRayMask &= (~cur->active)|move_mask(cmple(cur->exitT,quadRay.quadTEnd));
      if (!liveRayMask) return;
      continue;
    }
    //otherwise it must be a splitting node
    const F32vec4 splitT = (F32vec4(node->splitValue) - quadRay.quadOrigin[type])
                              * quadRay.quadInvDir[type];
//    const F32vec4 exitT = simd_min(cur->exitT,quadRay.quadTEnd); //ray's end may have changed since exitT was computed
    const F32vec4 exitT = cur->exitT;
    const F32vec4 enterT = cur->enterT;
    int leftIndex = quadRay.ray[0].dirType(type);
    const F32vec4 aExitT = simd_min(exitT,splitT);
    const F32vec4 bEnterT = simd_max(splitT,enterT);
//    cur[leftIndex].nodeHandle = node->left;
//    cur[1-leftIndex].nodeHandle = node->right;
    cur[1].active = move_mask(cmple(enterT,aExitT));
    cur[1].nodeHandle = node->getChild(1-leftIndex);
    cur[0].nodeHandle = node->getChild(leftIndex);
    cur[1].enterT = enterT;
    cur[1].exitT = aExitT;
    cur[0].active = move_mask(cmple(bEnterT,exitT));
    cur[0].enterT = bEnterT;
    cur[0].exitT = exitT;
    cur+=2;
  }
}

